/******************************************************************************
 * Copyright(c)  2005-2010, mqbest2000@gmail.com
 * 
 * Filename:	stun.c
 * Version:	1.0
 * Author:	MeiQin
 * Date:		13:05pm, 06.Dec.2010
 * Description: 

 * modification history
 * 		Version:	1.0
 * 		Author:	MeiQin
 * 		Date:	13:05pm, 06.Dec.2010
 * 		Modification:	
 ******************************************************************************/

#include "xsstun.h"

char *encode16(char* buf, UInt16 data)
{
	UInt16 ndata = htons(data);
	memcpy(buf, &ndata, sizeof(UInt16));
	return buf + sizeof(UInt16);
}

char *encode32(char* buf, UInt32 data)
{
	UInt32 ndata = htonl(data);
	memcpy(buf, &ndata, sizeof(UInt32));
	return buf + sizeof(UInt32);
}

UInt128 createTransctionID (void)
{
	UInt128 id;
	UInt32 i,r;
	static int iRand = 0;
	
	iRand++;
	srand(iRand);
	for(i=0; i<16; i+=4)
	{
		r = rand();
		id.octet[i+0] = r>>0;
		id.octet[i+1] = r>>8;
		id.octet[i+2] = r>>16;
		id.octet[i+3] = r>>24;
	}
	return id;
}

void initIdArray(UInt128 *ID, int length)
{
	int i,j;
	
	for(i=0; i<length; i++)
	{
		for(j=0; j<16; j++)
		{
			ID->octet[j] = 0;
		}
		ID++;
	}
}

int checkTransctionID(UInt128 id, UInt128 *ID, int count)
{
	int i;
	//int j;
	//xs_print("count is %d\n",count);
		
	for(i=0; i<count; i++)
	{
		//xs_print("\n");
		//for(j=0; j<16; j++)xs_print("%x",ID->octet[j]);
		//xs_print("\n");
		//for(j=0;j<16;j++)xs_print("%x",id.octet[j]);
		//xs_print("\n");
		if(memcmp(ID,&id,16) == 0)return _TRUE;
		else ID++;
	}

	return _FALSE;
}
	
void buildSimBindingRequest(unsigned char *buf, UInt128 *ID)
{	
	StunMsgHdr header;
	char *p;

	header.msgType = BindingRequest;
	header.msgLength = 0;
	header.id = createTransctionID();
	*ID = header.id;
	p = encode16((char *)buf, header.msgType);
	p = encode16(p, header.msgLength);
	memcpy(p, &header.id, sizeof(header.id));
}

void buildTestBindingRequest(unsigned char *buf, int change_ip, int change_port,
                              UInt128 *ID)
{
	StunMsgHdr header;
	StunAtrHdr atr_head;
	StunAtrChangeRequest atr;
	char *p;
	//write stun header
	header.msgType = BindingRequest;
	header.msgLength = 8;
	header.id = createTransctionID();
	*ID = header.id;
	p = encode16((char *)buf, header.msgType);
	p = encode16(p, header.msgLength);
	memcpy(p, &header.id, sizeof(header.id));
	p = p+sizeof(header.id);

	//write stun attributes
	atr_head.type = ChangeRequest;
	atr_head.length = 4;
	atr.value = 0x00000000;
	if(change_ip)
	{
		atr.value = atr.value|0x00000004;
	}
	if(change_port)
	{
		atr.value = atr.value|0x00000002;
	}
	atr.value = htonl(atr.value);
	p = encode16(p, atr_head.type);
	p = encode16(p, atr_head.length); 
	memcpy(p, &atr.value, sizeof(atr.value));
	//for test
	#if 0
	int i;
	for(i=0; i<28;i++)
	{
		xs_print("%x ",*(buf+i));
		if(i == 19)printf("\n");
	}
	xs_print("\n");
	#endif
}
	
/********************************************/
/* buf: point to Resquest you want to send  */
/* fd : file descriptor returned by socket  */
/* ip : address to send.(Network-endian)    */
/*port: port to send.(Network-endian)       */
/*  l : length of request you send          */
/********************************************/
int sendBindingRequest(xsCtx_t *pstCtx, unsigned char *buf, UInt16 len, UInt32 ip, UInt16 port)
{
	int nbyte;

	nbyte = pstCtx->pfUdpSend(ip, port, buf, len);
	if(nbyte == -1){
		xs_print("send failed\n");
		return _FALSE;
	}
	return _TRUE;
}

 /****************************************************/
 /* buf: point to Response you have received         */
 /* fd : file descriptor returned by socket          */
 /*  l : max length you can receive                  */
 /****************************************************/
int receiveResponse(xsCtx_t *pstCtx, unsigned char *buf, UInt16 len)
{
	int nbyte;

	nbyte = pstCtx->pfUdpRecv(0,0, buf, len);
	return nbyte;
}

//if ipv6,rewrite this function
int parseAtrAddress(const char *p, UInt16 length, StunAtrAddress4 *addr)
{	
	if(length != 8)return _FALSE;

	memcpy(&addr->ipv4.port, p+2, 2);
	memcpy(&addr->ipv4.addr, p+4, 4);
	addr->ipv4.port = ntohs(addr->ipv4.port);
	addr->ipv4.addr = ntohl(addr->ipv4.addr);
	
	return _TRUE;
}

/************************************************/
/*Unknown response return _FALSE                */
/*incorrect response return _FALSE               */
/*correct response rertun _TRUE                 */
/**********************************************/
int parseBindingResponse(unsigned char *buf, Address4 *addr, StunAddressInfo *stAddressinfo, UInt128 *ID, int count)
{	
	StunMessage msg;
	StunAtrHdr atr_head;
	char *p=(char *)buf;
	UInt8 *q = 0;
	int len;
	// init message
	msg.hasMappedAddress    = _FALSE;
	msg.hasSourceAddress    = _FALSE;
	msg.hasChangedAddress   = _FALSE;
	msg.hasMessageIntegrity = _FALSE;
	msg.hasReflectedFrom	= _FALSE;

	memcpy(&msg.msgHdr.msgType, p, 2);
	memcpy(&msg.msgHdr.msgLength, p+2, 2);
	memcpy(&msg.msgHdr.id, p+4, 16);
	if(!checkTransctionID(msg.msgHdr.id,ID,count))
	{
		return _FALSE;
	}
	p = p+20;
	msg.msgHdr.msgType = ntohs(msg.msgHdr.msgType);
	msg.msgHdr.msgLength = ntohs(msg.msgHdr.msgLength);
	len = msg.msgHdr.msgLength;
	#if 0
	xs_print("message type:%x\t length:%d\n", msg.msgHdr.msgType,len);
	#endif
	
	if(msg.msgHdr.msgType == BindingResponse)
	{	
		while(len >0)
		{
			memcpy(&atr_head.type, p, 2);
			memcpy(&atr_head.length, p+2, 2);
			atr_head.type = htons(atr_head.type);
			atr_head.length = htons(atr_head.length);
			p = p+4;
			switch(atr_head.type)
			{
				case MappedAddress:
					msg.hasMappedAddress = _TRUE;
					if(!parseAtrAddress(p,atr_head.length, &msg.mappedAddress))
					{
						xs_print("parse MappedAddress failed!\n\r");
						return _FALSE;
					}
					break;
				case SourceAddress:
					msg.hasSourceAddress = _TRUE;
					if(!parseAtrAddress(p,atr_head.length,
								&msg.sourceAddress))
					{
						xs_print("parse SourceAddress failed!\n\r");
						return _FALSE;
					}
					break;
				case ChangedAddress:
					msg.hasChangedAddress = _TRUE;
					if(!parseAtrAddress(p,atr_head.length,
								&msg.changedAddress))
					{
						xs_print("parse ChangedAddress failed!\n\r");
						return _FALSE;
				}
				break;
				case MessageIntegrity:break;
				case ReflectedFrom:break;
				default:
					xs_print("Unknown attribute:%x\n\r",atr_head.type);
					break;
			}
			len = len - 4 -atr_head.length;
			p = p+atr_head.length;
		}
		//output information to file
		memset(stAddressinfo, 0, sizeof(StunAddressInfo));
		if(msg.hasMappedAddress)
		{	
			addr->mapped_ip = msg.mappedAddress.ipv4.addr;
			addr->mapped_port = msg.mappedAddress.ipv4.port;
			q = (UInt8 *)&msg.mappedAddress.ipv4.addr;
			sprintf((char *)stAddressinfo->acMappedAddress,"%d.%d.%d.%d",*(q+3),*(q+2),*(q+1),*q);
			stAddressinfo->iMappePort = msg.mappedAddress.ipv4.port;
			//xs_print("MappedAddress:%s:%d\n\r",stAddressinfo->acMappedAddress, stAddressinfo->iMappePort);
		}
		else
		{
			xs_print("no MappedAddress attributes\n\r");
		}

		if(msg.hasSourceAddress)
		{
			q = (UInt8 *)&msg.sourceAddress.ipv4.addr;
			sprintf((char *)stAddressinfo->acSourceAddress,"%d.%d.%d.%d",*(q+3),*(q+2),*(q+1),*q);
			stAddressinfo->iSourcePort = msg.sourceAddress.ipv4.port;
			//xs_print("SourceAddress:%s:%d\n",stAddressinfo->acSourceAddress,stAddressinfo->iSourcePort);
		}
		else
		{
			xs_print("no SourceAddress attribute!\n\r");
		}

		if(msg.hasChangedAddress)
		{
			addr->changed_ip = msg.changedAddress.ipv4.addr;
			addr->changed_port = msg.changedAddress.ipv4.port;
			q =(UInt8 *)&msg.changedAddress.ipv4.addr;
			sprintf((char *)stAddressinfo->acChangedAddress,"%d.%d.%d.%d",*(q+3),*(q+2),*(q+1),*q);
			stAddressinfo->iChangePort= msg.changedAddress.ipv4.port;
			//xs_print("ChangedAddress:%s:%d\n",stAddressinfo->acChangedAddress, stAddressinfo->iChangePort);
		}
		else
		{
			xs_print("no ChangedAddress attribute!\n\r");
		}
		
		if(msg.hasMessageIntegrity)
		{
			xs_print("hasMessageIntegrity attribute!\n\r");
			//do something
		}
		
		if(msg.hasReflectedFrom)
		{
			xs_print("hasReflectedFrom attribute!\n\r");
			//do something
		}
	}//end of parse binding response
	
	else if(msg.msgHdr.msgType == BindingErrorResponse)
	{
		xs_print("error code\n\r");		
	}
	else
	{
		xs_print("Unknown stun type!\n\r");
		return _FALSE;
	}
	xs_print("\n\r");

	return _TRUE;
}

int GetMapAddr(xsCtx_t *pstCtx, StunAddressInfo *pstAddressinfo, unsigned int local_ip, unsigned int dst_ip, unsigned short dst_p)
{
	int iRst = XS_CONTINUE;
	static int iStep = 0;
	static int state = Build;
	static int count = 0;
	static int change_ip = _FALSE;
	static int change_port = _FALSE;
	static unsigned short usPort = 0;
	static Address4 stAddr;
	static UInt128 ID[10];

	unsigned char bufReq[REQ_LEN] = {0};
	unsigned char bufRes[RES_LEN] = {0};

	if(0 == count){
		initIdArray(ID,10);	
	}
	xs_print("GetMapAddr->Step:%d, State:%d, count:%d\n\r", iStep, state, count);
	switch(state){
		case Build:
			buildTestBindingRequest(bufReq,change_ip,change_port,&ID[count]);
			if(3 == iStep){
				dst_ip = htonl(stAddr.changed_ip);
				dst_p = htons(stAddr.changed_port);
			}
			sendBindingRequest(pstCtx, bufReq, REQ_LEN,dst_ip,dst_p);
			count++;
			state = WaitResponse;	
			break;
		case WaitResponse:
			iRst = receiveResponse(pstCtx, bufRes, RES_LEN);
			
			if(0 < iRst){
				if(parseBindingResponse(bufRes,&stAddr,pstAddressinfo, ID, count)){
					state = EndProcess;
				}
				iRst = XS_CONTINUE;
			}else if(0 == iRst){ // timeout
				if(count < TYR_TIMES){
					state = Build;
				}
				iRst = XS_CONTINUE;
			}else{
				iRst = XS_FAILURE;
			}
			break;
		case EndProcess:
			break;
	}
	
	switch(iStep){
		case 0:
			if(EndProcess == state){
				xs_print("LocalIP:%X, MapIp:%X\n\r", local_ip, stAddr.mapped_ip);
				if(local_ip == stAddr.mapped_ip){
					iStep = 1;
				}else{
					iStep = 2;
				}
				usPort = stAddr.mapped_port;
				change_ip = _TRUE;
				change_port = _TRUE;
				state = Build;
				count = 0;
			}
			break;
		case 1:
			if(EndProcess == state){
				pstAddressinfo->cNatType = 1; // open
			}else if(count >= TYR_TIMES){
				pstAddressinfo->cNatType = 12; //close
			} 
			iRst = XS_OK;
			break;
		case 2:
			if(EndProcess == state){
				pstAddressinfo->cNatType = 3; //full
				iRst = XS_OK;
			}else if(count >= TYR_TIMES){
				iStep = 3;
				change_ip = _FALSE;
				change_port = _FALSE;
				state = Build;
				count = 0;
			}
			break;
		case 3:
			if(count >= TYR_TIMES || EndProcess == state){
				if(usPort == stAddr.mapped_port){
					iStep = 4;
					change_ip = _FALSE;
					change_port = _TRUE;
					state = Build;
					count = 0;
				}else{
					pstAddressinfo->cNatType = 6; //symmetric
					iRst = XS_OK;
				}
			}
			break;
		case 4:
			if(EndProcess == state){
				pstAddressinfo->cNatType = 4; // restricted
				iRst = XS_OK;
			}else if(count >= TYR_TIMES){
				pstAddressinfo->cNatType = 5; //port_restricted
				iRst = XS_OK;
			} 
			break;
	}
	
	if(XS_CONTINUE != iRst){
		iStep = 0;
		state = Build;
		count = 0;
    change_ip = _FALSE;
    change_port = _FALSE;
	}
	return iRst;
}


